import pygame, sys
from dataclasses import dataclass
from enum import IntEnum, Enum
import numpy as np


class E_MapType(IntEnum):
    Blank: int = 0
    Wall: int = 1
    StartPoint: int = 2
    EndPoint: int = 3


class E_MapColor(Enum):
    """
    地图不同标记物颜色, 颜色顺序RGB
    """

    Gap = (255, 255, 255)  # 间隔-白色
    Blank = (0, 255, 0)  # 空白-绿色
    Wall = (0, 0, 0)  # 墙-黑色
    StartPoint = (0, 0, 255)  # 起点-蓝色
    EndPoint = (255, 0, 0)  # 终点-红色


@dataclass
class MapData:
    w: int = 5
    h: int = 5
    screenScale: int = 100
    lineSize: int = 5
    data: any = None


class gameRect:
    def __init__(self, mapData: MapData):
        pygame.init()
        self.map = mapData
        self.screenScale = self.map.screenScale
        self.width = self.map.w * self.screenScale
        self.height = self.map.h * self.screenScale
        self.screen = pygame.display.set_mode((self.width, self.height))
        self.clock = pygame.time.Clock()
        self.init_map()
        self.init_font()
        self.init_data()

    def __exit__(self):
        pygame.quit()

    def init_map(self):
        """
        初始化地图数据
        """
        self.map.data = np.zeros(self.map.w * self.map.h, dtype=np.int32).reshape((self.map.w, self.map.h))

    def init_font(self):
        """初始化字体"""
        font_path = "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc"
        font_size = 24
        self.font = pygame.font.Font(font_path, font_size)

    def init_data(self):
        self.mapType = E_MapType.Wall
        self.mouse_pos = (0, 0)

    def draw_map(self):
        """
        绘制网格线
        """
        lineSize = self.map.lineSize
        # 遍历行 * 列
        for i in range(self.map.w):
            for j in range(self.map.h):
                # 计算当前方格左上角坐标
                x = i * self.screenScale + lineSize
                y = j * self.screenScale + lineSize
                map_type = self.map.data[i][j]
                if map_type == E_MapType.Blank.value:
                    color = E_MapColor.Blank.value
                    pygame.draw.rect(
                        self.screen, color, (x, y, self.screenScale - lineSize, self.screenScale - lineSize), 0
                    )
                elif map_type == E_MapType.Wall.value:
                    color = E_MapColor.Wall.value
                    pygame.draw.rect(
                        self.screen, color, (x, y, self.screenScale - lineSize, self.screenScale - lineSize), 0
                    )
                elif map_type == E_MapType.StartPoint.value:
                    color = E_MapColor.StartPoint.value
                    pygame.draw.rect(
                        self.screen, color, (x, y, self.screenScale - lineSize, self.screenScale - lineSize), 0
                    )
                    self.draw_text("start", (x, y + self.screenScale // 3), (255, 255, 255))
                elif map_type == E_MapType.EndPoint.value:
                    color = E_MapColor.EndPoint.value
                    pygame.draw.rect(
                        self.screen, color, (x, y, self.screenScale - lineSize, self.screenScale - lineSize), 0
                    )
                    self.draw_text("end", (x, y + self.screenScale // 3), (255, 255, 255))
                else:
                    continue

    def draw_note_text(self):
        """
        绘制提示文字: wall,start,end
        """
        pos = self.mouse_pos
        color = (255, 255, 255)
        if self.mapType == E_MapType.Wall:
            text = "wall"
        elif self.mapType == E_MapType.StartPoint:
            text = "start"
        elif self.mapType == E_MapType.EndPoint:
            text = "end"
        self.draw_text(text, pos, color)

    def draw_text(self, text, pos, color):
        text_surface = self.font.render(text, True, color)
        self.screen.blit(text_surface, pos)

    def get_map_idx_by_screen_pos(self, pos):
        """
        根据屏幕坐标获取地图坐标
        """
        x = pos[0] // self.screenScale
        y = pos[1] // self.screenScale
        return x, y

    def handler_mouse_event(self, event):
        """
        鼠标按下事件
            左键: 标记wall, start, end
            右键: 标记空白
        """
        map_idx = self.get_map_idx_by_screen_pos(event.pos)
        if event.button == 1:  # 左键
            map_type = self.mapType.value
            self.map.data[map_idx[0]][map_idx[1]] = map_type
        elif event.button == 3:
            self.map.data[map_idx[0]][map_idx[1]] = E_MapType.Blank.value

    def handler_mouse_motion(self, event):
        """
        鼠标移动事件
        """
        self.mouse_pos = event.pos

    def handler_key_press(self, event):
        """
        按键按下
            1 标记类型wall
            2 标记类型start
            3 标记类型end
        """
        if event.key == pygame.K_1:
            self.mapType = E_MapType.Wall
        elif event.key == pygame.K_2:
            self.mapType = E_MapType.StartPoint
        elif event.key == pygame.K_3:
            self.mapType = E_MapType.EndPoint

    def loop_run(self):
        self.color_idx = 0
        while True:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    sys.exit()
                elif event.type == pygame.MOUSEBUTTONDOWN:
                    self.handler_mouse_event(event)
                elif event.type == pygame.MOUSEMOTION:
                    self.handler_mouse_motion(event)
                elif event.type == pygame.KEYDOWN:
                    self.handler_key_press(event)

            color = E_MapColor.Gap.value
            self.screen.fill(color)  # 清屏
            self.draw_map()  # 绘制地图
            self.draw_note_text()

            pygame.display.update()  # 刷新屏幕
            self.clock.tick(1000)  # 限制fps 1帧


if __name__ == "__main__":
    map_data = MapData()
    gameImpl = gameRect(map_data)
    gameImpl.loop_run()
